"""
AI执行协调器
"""
import asyncio
import logging
from datetime import datetime
from typing import List, Dict, Any, Optional, Callable

from backend.services.ai_service import get_ai_service
from backend.services.question_analyzer import simple_question_handler
from backend.models.core import (
    Task, TaskResult, ExecutionSession, ExecutionResult, 
    Question, QuestionAnalysis
)
from backend.models.session_manager import session_manager
from shared.enums import (
    TaskStatus, SessionStatus, ExecutionPath, 
    ComplexityLevel, TaskType
)

logger = logging.getLogger(__name__)


class AICoordinator:
    """AI执行协调器 - 调度和执行任务"""
    
    def __init__(self):
        self.task_execution_prompts = {
            TaskType.RESEARCH: """
请执行以下研究任务：

任务描述：{description}

相关背景信息：
{context}

请进行深入研究，收集相关信息和数据。提供准确、全面的研究结果。
""",
            TaskType.ANALYSIS: """
请执行以下分析任务：

任务描述：{description}

相关信息：
{context}

请进行深入分析，提供清晰的分析结果和结论。
""",
            TaskType.SYNTHESIS: """
请执行以下综合任务：

任务描述：{description}

需要综合的信息：
{context}

请将相关信息进行综合整理，形成完整的结论。
""",
            TaskType.GENERATION: """
请执行以下生成任务：

任务描述：{description}

参考信息：
{context}

请根据要求生成相应的内容。
""",
            TaskType.EVALUATION: """
请执行以下评估任务：

任务描述：{description}

评估对象：
{context}

请进行客观公正的评估，提供评估结果和理由。
""",
            TaskType.COMPARISON: """
请执行以下比较任务：

任务描述：{description}

比较对象：
{context}

请进行详细比较，指出异同点和各自特点。
"""
        }
        
        # 事件回调
        self.event_callbacks: Dict[str, List[Callable]] = {
            'execution_started': [],
            'task_started': [],
            'task_completed': [],
            'execution_completed': [],
            'execution_failed': []
        }
    
    def add_event_callback(self, event: str, callback: Callable):
        """添加事件回调"""
        if event in self.event_callbacks:
            self.event_callbacks[event].append(callback)
    
    async def _emit_event(self, event: str, data: Dict[str, Any]):
        """触发事件"""
        if event in self.event_callbacks:
            for callback in self.event_callbacks[event]:
                try:
                    if asyncio.iscoroutinefunction(callback):
                        await callback(data)
                    else:
                        callback(data)
                except Exception as e:
                    logger.error(f"Event callback error: {e}")
    
    async def execute_simple_question(self, question: str) -> ExecutionResult:
        """执行简单问题"""
        try:
            logger.info(f"Executing simple question: {question[:50]}...")
            
            result = await simple_question_handler.handle_simple_question(question)
            
            if result.success:
                logger.info("Simple question executed successfully")
            else:
                logger.error(f"Simple question execution failed: {result.error}")
            
            return result
            
        except Exception as e:
            logger.error(f"Simple question execution error: {e}")
            return ExecutionResult(
                success=False,
                error=str(e)
            )
    
    async def execute_task_sequence(self, session_id: str) -> ExecutionResult:
        """执行任务序列"""
        try:
            session = session_manager.get_session(session_id)
            if not session:
                return ExecutionResult(
                    success=False,
                    error="Session not found"
                )
            
            logger.info(f"Starting task sequence execution for session: {session_id}")
            
            # 更新会话状态
            session.status = SessionStatus.EXECUTING
            session.started_at = datetime.now()
            session_manager.update_session(session)
            
            # 触发执行开始事件
            await self._emit_event('execution_started', {
                'session_id': session_id,
                'task_count': len(session.tasks)
            })
            
            # 执行任务序列
            execution_context = {}
            total_tokens = 0
            
            while True:
                # 获取准备执行的任务
                ready_tasks = session.get_ready_tasks()
                
                if not ready_tasks:
                    # 检查是否所有任务都完成
                    completed_tasks = [task for task in session.tasks if task.status == TaskStatus.COMPLETED]
                    if len(completed_tasks) == len(session.tasks):
                        break  # 所有任务完成
                    
                    # 检查是否有失败的任务
                    failed_tasks = [task for task in session.tasks if task.status == TaskStatus.FAILED]
                    if failed_tasks:
                        logger.error(f"Found failed tasks: {[task.id for task in failed_tasks]}")
                        break
                    
                    # 可能存在循环依赖或其他问题
                    logger.warning("No ready tasks found but not all tasks completed")
                    break
                
                # 执行准备好的任务（可以并行执行无依赖的任务）
                task_results = await self._execute_tasks_batch(ready_tasks, execution_context, session_id)
                
                # 更新会话中的任务状态和结果
                for task_result in task_results:
                    task = session.get_task_by_id(task_result.task_id)
                    if task:
                        if task_result.status == TaskStatus.COMPLETED:
                            task.mark_completed(task_result.result or "")
                            execution_context[task.id] = task_result.result
                        else:
                            task.mark_failed(task_result.error or "Execution failed")
                        
                        if task_result.tokens_used:
                            total_tokens += task_result.tokens_used
                
                # 添加结果到会话
                session.results.extend(task_results)
                session.total_tokens_used = total_tokens
                session_manager.update_session(session)
            
            # 生成最终答案
            final_answer = await self._generate_final_answer(session, execution_context)
            
            # 更新会话完成状态
            session.final_answer = final_answer
            session.status = SessionStatus.COMPLETED
            session.completed_at = datetime.now()
            session.total_tokens_used = total_tokens
            session_manager.update_session(session)
            
            # 触发执行完成事件
            await self._emit_event('execution_completed', {
                'session_id': session_id,
                'final_answer': final_answer,
                'tokens_used': total_tokens
            })
            
            logger.info(f"Task sequence execution completed for session: {session_id}")
            
            return ExecutionResult(
                success=True,
                result=final_answer,
                tokens_used=total_tokens,
                execution_time=(session.completed_at - session.started_at).total_seconds() if session.completed_at and session.started_at else 0
            )
            
        except Exception as e:
            logger.error(f"Task sequence execution error: {e}")
            
            # 更新会话失败状态
            session = session_manager.get_session(session_id)
            if session:
                session.status = SessionStatus.FAILED
                session.completed_at = datetime.now()
                session_manager.update_session(session)
            
            # 触发执行失败事件
            await self._emit_event('execution_failed', {
                'session_id': session_id,
                'error': str(e)
            })
            
            return ExecutionResult(
                success=False,
                error=str(e)
            )
    
    async def _execute_tasks_batch(self, tasks: List[Task], context: Dict[str, Any], session_id: str) -> List[TaskResult]:
        """批量执行任务"""
        if not tasks:
            return []
        
        # 并行执行任务
        task_coroutines = [self._execute_single_task(task, context, session_id) for task in tasks]
        results = await asyncio.gather(*task_coroutines, return_exceptions=True)
        
        task_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                logger.error(f"Task execution failed: {result}")
                task_results.append(TaskResult(
                    task_id=tasks[i].id,
                    status=TaskStatus.FAILED,
                    error=str(result)
                ))
            else:
                task_results.append(result)
        
        return task_results
    
    async def _execute_single_task(self, task: Task, context: Dict[str, Any], session_id: str) -> TaskResult:
        """执行单个任务"""
        start_time = datetime.now()
        
        try:
            logger.info(f"Executing task: {task.id} - {task.description[:50]}...")
            
            # 触发任务开始事件
            await self._emit_event('task_started', {
                'session_id': session_id,
                'task_id': task.id,
                'task_description': task.description
            })
            
            # 标记任务开始
            task.mark_started()
            
            # 准备上下文信息
            task_context = self._build_task_context(task, context)
            
            # 获取任务执行提示
            prompt_template = self.task_execution_prompts.get(task.type, self.task_execution_prompts[TaskType.RESEARCH])
            prompt = prompt_template.format(
                description=task.description,
                context=task_context
            )
            
            # 执行任务
            result = await get_ai_service().generate_text(
                prompt,
                temperature=0.6,
                max_tokens=2000
            )
            
            execution_time = (datetime.now() - start_time).total_seconds()
            
            if result.success:
                logger.info(f"Task completed: {task.id}")
                
                task_result = TaskResult(
                    task_id=task.id,
                    status=TaskStatus.COMPLETED,
                    result=result.result,
                    execution_time=execution_time,
                    tokens_used=result.tokens_used
                )
                
                # 触发任务完成事件
                await self._emit_event('task_completed', {
                    'session_id': session_id,
                    'task_id': task.id,
                    'result': result.result,
                    'execution_time': execution_time
                })
                
                return task_result
            else:
                logger.error(f"Task failed: {task.id} - {result.error}")
                return TaskResult(
                    task_id=task.id,
                    status=TaskStatus.FAILED,
                    error=result.error,
                    execution_time=execution_time
                )
                
        except Exception as e:
            execution_time = (datetime.now() - start_time).total_seconds()
            logger.error(f"Task execution error: {task.id} - {e}")
            return TaskResult(
                task_id=task.id,
                status=TaskStatus.FAILED,
                error=str(e),
                execution_time=execution_time
            )
    
    def _build_task_context(self, task: Task, context: Dict[str, Any]) -> str:
        """构建任务执行上下文"""
        context_parts = []
        
        # 添加依赖任务的结果
        for dep_id in task.dependencies:
            if dep_id in context:
                context_parts.append(f"依赖任务结果：{context[dep_id]}")
        
        # 如果没有上下文，添加默认说明
        if not context_parts:
            context_parts.append("这是第一个任务，没有前置上下文。")
        
        return "\n\n".join(context_parts)
    
    async def _generate_final_answer(self, session: ExecutionSession, context: Dict[str, Any]) -> str:
        """生成最终答案"""
        try:
            # 收集所有成功任务的结果
            successful_results = []
            for task in session.tasks:
                if task.status == TaskStatus.COMPLETED and task.result:
                    successful_results.append(f"任务：{task.description}\n结果：{task.result}")
            
            if not successful_results:
                return "抱歉，没有任务成功完成，无法生成答案。"
            
            # 生成最终答案的提示
            final_prompt = f"""
基于以下任务执行结果，为用户问题生成一个完整、准确的最终答案。

原始问题：{session.question.content if session.question else '未知问题'}

任务执行结果：
{chr(10).join(successful_results)}

请整合所有结果，生成一个清晰、完整的答案。答案应该：
1. 直接回答用户的问题
2. 逻辑清晰，结构完整
3. 包含必要的细节和解释
4. 语言流畅自然

只返回最终答案，不要包含任务分解过程。
"""
            
            result = await get_ai_service().generate_text(
                final_prompt,
                temperature=0.5,
                max_tokens=3000
            )
            
            if result.success:
                return result.result
            else:
                logger.error(f"Final answer generation failed: {result.error}")
                return "基于执行结果整合答案时出现问题，但以下是各任务的结果：\n\n" + "\n\n".join(successful_results)
                
        except Exception as e:
            logger.error(f"Final answer generation error: {e}")
            return "生成最终答案时出现错误，请稍后再试。"
    
    async def monitor_execution(self, session_id: str) -> Dict[str, Any]:
        """监控执行状态"""
        session = session_manager.get_session(session_id)
        if not session:
            return {"error": "Session not found"}
        
        return {
            "session_id": session_id,
            "status": session.status.value,
            "progress": session.get_progress(),
            "total_tasks": len(session.tasks),
            "completed_tasks": len([task for task in session.tasks if task.status == TaskStatus.COMPLETED]),
            "failed_tasks": len([task for task in session.tasks if task.status == TaskStatus.FAILED]),
            "current_tasks": [
                {
                    "id": task.id,
                    "description": task.description,
                    "status": task.status.value
                }
                for task in session.tasks if task.status == TaskStatus.RUNNING
            ]
        }


# 全局协调器实例
ai_coordinator = AICoordinator()